const std = @import("std");

pub const FlagError = error{
    NoProgramName,
    BadFlagSynatx,
    NoArgsForFlag,
    InvalidValueForFlag,
    NoFlagForArg,
    ErrHelp,
};

const Value = union(enum) {
    bool_value: *bool,
    string_value: *[]const u8,
    i32_value: *i32,
    i64_value: *i64,
    isize_value: *isize,
    u32_value: *u32,
    u64_value: *u64,
    usize_value: *usize,

    fn set(self: *Value, value: []const u8) !void {
        switch (self.*) {
            .bool_value => |*flag_value| {
                flag_value.*.* = std.mem.eql(u8, value, "true");
            },
            .string_value => |*str_value| {
                str_value.*.* = value;
            },
            .i32_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(i32, value, 10);
            },
            .i64_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(i64, value, 10);
            },
            .isize_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(isize, value, 10);
            },
            .u32_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(u32, value, 10);
            },
            .u64_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(u64, value, 10);
            },
            .usize_value => |*int_value| {
                int_value.*.* = try std.fmt.parseInt(usize, value, 10);
            },
        }
    }

    fn toString(self: Value, allocator: std.mem.Allocator) ![]u8 {
        switch (self) {
            .string_value => |value| {
                return allocator.dupe(u8, value.*);
            },
            .bool_value => |value| {
                return allocator.dupe(u8, if (value.*) "true" else "false");
            },
            inline else => |value| {
                return std.fmt.allocPrint(allocator, "{d}", .{value.*});
            },
        }
    }

    // TODO
    pub fn get(comptime T: type, value: *T) Value {
        switch (T) {
            bool => return Value{ .bool_value = value },
            []const u8 => return Value{ .string_value = value },
            i32 => return Value{ .i32_value = value },
            i64 => return Value{ .i64_value = value },
            isize => return Value{ .isize_value = value },
            u32 => return Value{ .u32_value = value },
            u64 => return Value{ .u64_value = value },
            usize => return Value{ .usize_value = value },
            else => @compileError("Value.get() 不支持的类型: " ++ @typeName(T)),
        }
    }
};

pub const ErrorHandling = enum(u8) {
    ContinueOnError,
    ExitOnError,
    PanicOnError,
};

pub const Flag = struct {
    allocator: std.mem.Allocator,
    name: []const u8,
    usage: []const u8,
    value: Value,
    def_value: []u8,

    pub fn deinit(self: *Flag) void {
        self.allocator.free(self.def_value);
    }
};

const StringFlagHashMap = std.StringHashMap(Flag);

pub const UsageFn = *const fn (*FlagSet, *std.io.Writer) anyerror!void;

pub const FlagSet = struct {
    allocator: std.mem.Allocator,
    name: []const u8,
    is_parsed: bool,
    flags: std.ArrayList(Flag),
    args: std.ArrayList([]const u8) = .empty,
    actual: StringFlagHashMap,
    formal: StringFlagHashMap,
    args_iter: std.process.ArgIterator,
    error_handling: ErrorHandling,
    usage_fn: UsageFn = printDefault,
    output: *?std.io.Writer = null,

    pub fn setOutput(self: *FlagSet, writer: *std.io.Writer) void {
        self.output = writer;
    }

    pub fn usage(self: *FlagSet) !void {
        if (self.output == null) {
            var stdout_buffer: [1024 * 10]u8 = undefined;
            var stdout_writer = std.fs.File.stdout().writer(&stdout_buffer);
            var stdout = &stdout_writer.interface;
            try self.usageWithWriter(stdout);
            try stdout.flush();
        } else {
            try self.usageWithWriter(self.output.?);
            try self.output.?.flush();
        }
    }

    fn usageWithWriter(self: *FlagSet, writer: *std.io.Writer) !void {
        try self.usage_fn(self, writer);
    }

    pub fn setUsage(self: *FlagSet, custom_flag: UsageFn) void {
        self.usage_fn = custom_flag;
    }

    pub fn deinit(self: *FlagSet) void {
        for (self.flags.items) |*item| {
            item.deinit();
        }
        self.flags.deinit(self.allocator);
        self.actual.deinit();
        self.formal.deinit();
        self.args_iter.deinit();
        self.args.deinit(self.allocator);
    }

    pub fn flagVar(self: *FlagSet, comptime T: type, value: *T, comptime name: []const u8, def_val: T, usage_arg: []const u8) !void {
        comptime {
            if (std.mem.startsWith(u8, name, "-")) {
                @compileError(std.fmt.comptimePrint("flag {s} begins with -", .{name}));
            } else if (std.mem.indexOf(u8, name, "=") != null) {
                @compileError(std.fmt.comptimePrint("flag {s} contains =", .{name}));
            }
        }

        value.* = def_val;
        const flag_value = Value.get(T, value);

        if (self.formal.get(name) != null) {
            const msg = if (self.name.len == 0)
                try std.fmt.allocPrint(self.allocator, "flag redefined: {s}", .{name})
            else
                try std.fmt.allocPrint(self.allocator, "{s} flag redefined: {s}", .{ self.name, name });
            @panic(msg);
        }
        const flag = try self.flags.addOne(self.allocator);
        flag.* = Flag{
            .allocator = self.allocator,
            .name = name,
            .value = flag_value,
            .usage = usage_arg,
            .def_value = try flag_value.toString(self.allocator),
        };

        try self.formal.put(name, flag.*);
    }

    pub fn getArgs(self: *FlagSet) []const []const u8 {
        return self.args.items;
    }

    pub fn parsed(self: *FlagSet) !void {
        self.is_parsed = true;
        while (true) {
            const seen = self.parsedOne() catch |err| {
                switch (self.error_handling) {
                    .ContinueOnError => {
                        return err;
                    },
                    .ExitOnError => {
                        if (err == error.ErrHelp) {
                            std.process.exit(0);
                        }
                        std.process.exit(2);
                    },
                    .PanicOnError => {
                        @panic(try std.fmt.allocPrint(self.allocator, "error: {any}", .{err}));
                    },
                }
            };
            if (!seen) {
                break;
            }
        }
        while (self.args_iter.next()) |arg| {
            try self.args.append(self.allocator, arg);
        }
    }

    fn parsedOne(self: *FlagSet) !bool {
        const arg = self.args_iter.next() orelse return false;
        if (arg.len < 2 or arg[0] != '-') {
            try self.args.append(self.allocator, arg);
            return false;
        }
        var numMinuses: usize = 1;
        if (arg[1] == '-') {
            numMinuses += 1;
            if (arg.len == 2) {
                try self.args.append(self.allocator, arg);
                return false;
            }
        }
        var name = arg[numMinuses..arg.len];
        if (name.len == 0 or name[0] == '-' or name[0] == '=') {
            try self.faildf("bad flag syntax: {s}", .{arg});
            return error.BadFlagSynatx;
        }
        var has_value = false;
        var value: ?[]const u8 = null;
        for (name, 0..) |c, i| {
            if (c == '=') {
                value = name[i + 1 ..];
                has_value = true;
                name = name[0..i];
                break;
            }
        }

        var flag = self.formal.get(name);
        if (flag == null) {
            if (std.mem.eql(u8, name, "help") or std.mem.eql(u8, name, "h")) {
                try self.usage();
                return error.ErrHelp;
            }
            try self.faildf("flag provided but not defined: -{s}\n", .{ name });
            return error.NoFlagForArg;
        }
        switch (flag.?.value) {
            .bool_value => |flag_value| {
                if (has_value) {
                    try flag.?.value.set(value.?);
                } else {
                    flag_value.* = true;
                }
            },
            inline else => |_| {
                if (!has_value) {
                    // value is the next arg
                    value = self.args_iter.next();
                    if (value == null) {
                        try self.faildf( "flag needs an argument: {s}", .{name});
                        return error.NoArgsForFlag;
                    }
                    flag.?.value.set(value.?) catch |e| {
                        try self.faildf("invalid value {s} for flag {s}: {any}", .{ value.?, name, e });
                        return error.InvalidValueForFlag;
                    };
                }
            },
        }
        try self.actual.put(name, flag.?);
        return true;
    }

    fn faildf(self: *FlagSet, comptime fmt: []const u8, args: anytype) !void {
        std.debug.print(fmt, args);
        try self.usage();
    }

    pub fn visitAll(self: FlagSet, visit: fn (*Flag) void) !void {
        const sorted_flags = try self.sortFlags(self.formal);
        for (sorted_flags) |flag| {
            visit(flag);
        }
        self.allocator.free(sorted_flags);
    }

    fn sortFlags(self: FlagSet, flags: StringFlagHashMap) ![]*Flag {
        var sorted_flags = try self.allocator.alloc(*Flag, flags.count());
        var i: usize = 0;
        var iter = flags.iterator();
        while (iter.next()) |kv| {
            sorted_flags[i] = kv.value_ptr;
            i += 1;
        }
        std.sort.pdq(*Flag, sorted_flags, {}, sortFlagsCmp);
        return sorted_flags;
    }
};

fn printDefault(self: *FlagSet, writer: *std.io.Writer) anyerror!void {
    try writer.print("Usage: {s} [flags]\n", .{self.name});
    const sort_flags = try self.sortFlags(self.formal);
    for (sort_flags) |flag| {
        try printFlagDefault(self.allocator, writer, flag);
    }
    self.allocator.free(sort_flags);
}

pub fn sortFlagsCmp(_: void, a: *Flag, b: *Flag) bool {
    return std.mem.lessThan(u8, a.name, b.name);
}

fn printFlagDefault(gpa: std.mem.Allocator, writer: *std.io.Writer, flag: *Flag) !void {
    try writer.print("  -{s}", .{flag.name});
    const name_usage = try unquoteUsage(gpa, flag);
    const name = name_usage[0];
    const usage_var = name_usage[1];
    defer {
        // 只有当 usage_var 不等于原始 flag.usage 时才释放内存
        if (usage_var.ptr != flag.usage.ptr) {
            gpa.free(usage_var);
        }
    }

    if (name.len > 0) {
        try writer.writeByte(' ');
        _ = try writer.write(name);
    }
    // Boolean flag of one ASCII letter are so common we
    // treat them specially, putting their usage on the same line.
    if (writer.end <= 4) {
        try writer.writeByte('\t');
    } else {
        // Four spaces before the tab triggers good alignment
        // for both 4- and 8-space tab stops.
        _ = try writer.write("\n    \t");
    }
    const usage_buffer = try std.mem.replaceOwned(u8, gpa, usage_var, "\n", "\n    \t");
    _ = try writer.write(usage_buffer);
    gpa.free(usage_buffer);

    // Print the default value only if it differs to the zero value
    // for this flag type
    if (!(try isZeroValue(flag, flag.def_value))) {
        try writer.print("  (default {s})", .{flag.def_value});
    }
    try writer.writeByte('\n');
}

fn isZeroValue(flag: *Flag, value: []const u8) !bool {
    switch (flag.value) {
        .bool_value => {
            return std.mem.eql(u8, value, "false");
        },
        .i32_value, .i64_value, .isize_value, .u32_value, .u64_value, .usize_value => {
            return std.mem.eql(u8, value, "0");
        },
        .string_value => {
            return value.len == 0;
        },
    }
}

fn unquoteUsage(allocator: std.mem.Allocator, flag: *Flag) ![2][]const u8 {
    var name: []const u8 = undefined;
    // Look for a back-quoted name, but avoid the strings package.
    const usage_var = flag.usage;
    for (usage_var, 0..) |c, i| {
        if (c == '`') {
            for (i + 1..usage_var.len) |j| {
                if (usage_var[j] == '`') {
                    name = usage_var[i + 1 .. j];
                    // 创建一个新的字符串，去掉反引号
                    const new_usage = try std.mem.concat(allocator, u8, &[_][]const u8{ usage_var[0..i], name, usage_var[j + 1 .. usage_var.len] });
                    return [2][]const u8{ name, new_usage };
                }
            }
            break;
        }
    }
    // No explicit name, so use type if we can find one.
    name = "value";
    switch (flag.value) {
        .bool_value => {
            name = "";
        },
        .i32_value, .i64_value, .isize_value => {
            name = "int";
        },
        .u32_value, .u64_value, .usize_value => {
            name = "uint";
        },
        .string_value => {
            name = "string";
        },
    }
    return [2][]const u8{ name, usage_var };
}

pub fn init(allocator: std.mem.Allocator) !FlagSet {
    return try initWithErrorHandling(allocator, ErrorHandling.ExitOnError);
}

pub fn initWithErrorHandling(allocator: std.mem.Allocator, error_handling: ErrorHandling) !FlagSet {
    var arg_iter = try std.process.argsWithAllocator(allocator);
    const name = arg_iter.next() orelse @panic("No program name");
    return FlagSet{
        .allocator = allocator,
        .name = name,
        .is_parsed = false,
        .flags = try std.ArrayList(Flag).initCapacity(allocator, 4),
        .actual = StringFlagHashMap.init(allocator),
        .formal = StringFlagHashMap.init(allocator),
        .args_iter = arg_iter,
        .error_handling = error_handling,
    };
}
